# advanced_mcp_agent.py
import asyncio
from typing import Dict, List, Any
import ollama

class AdvancedMCPAgent:
    """高级MCP Agent支持异步和复杂工具链"""
    
    def __init__(self, model: str = "deepseek-r1:14b"):
        self.model = model
        self.conversation_history = []
        self.tool_registry = {}
        
    def register_tool(self, name: str, description: str, function: callable):
        """注册MCP工具"""
        self.tool_registry[name] = {
            'description': description,
            'function': function
        }
    
    async def process_with_planning(self, user_input: str) -> Dict[str, Any]:
        """使用规划机制处理复杂查询"""
        
        # 第一步：任务规划
        planning_prompt = f"""
用户查询: {user_input}

请分析这个任务是否需要多个步骤完成，并制定执行计划。
如果需要多个工具，请说明执行顺序。
"""
        plan = await self._async_chat(planning_prompt)
        
        # 第二步：逐步执行
        execution_prompt = f"""
任务计划: {plan}
用户查询: {user_input}

开始执行第一步:
"""
        return await self._execute_with_tools(execution_prompt)
    
    async def _async_chat(self, prompt: str) -> str:
        """异步调用Ollama"""
        response = ollama.chat(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            options={'temperature': 0.2}
        )
        return response['message']['content']